'''
详细见

https://percent4.github.io/NLP%EF%BC%88%E4%B9%9D%E5%8D%81%E5%9B%9B%EF%BC%89transformers%E6%A8%A1%E5%9D%97%E4%B8%AD%E7%9A%84DataCollator/

'''

from transformers import DefaultDataCollator

# data = [
#     {
#         "input_ids": [101, 2003, 1037, 102],  # 句子："I like [SEP]"
#         "attention_mask": [1, 1, 1],
#         "labels": [0, 1, 2]  # 每个 token 的标签（0：O，1：B-PER，2：I-PER）
#     },
#     {
#         "input_ids": [101, 3000, 102],  # 句子："He [SEP]"
#         "attention_mask": [1, 1, 1],
#         "labels": [0, 3, 0]  # 标签（3：B-LOC）
#     }
# ]



# collator = DefaultDataCollator()
# batch = collator(data)

# print(batch["input_ids"])
# # 输出（最长长度为 4，短句子补 0）：
# # tensor([[101, 2003, 1037, 102],
# #         [101, 3000, 102,   0]])

# print(batch["labels"])
# # 输出（直接复制标签，短句子补 0，这会导致损失计算错误）：
# # tensor([[0, 1, 2, 0],
# #         [0, 3, 0, 0]])  # 最后一个 0 是无效标签，但会被计入损失



# # 文本分类数据（input_ids 长度不同，但 labels 都是单值）
# data = [
#     {"input_ids": [101, 2003, 102], "labels": 0},  # 长度 3
#     {"input_ids": [101, 3000, 4000, 102], "labels": 1}  # 长度 4
# ]

# collator = DefaultDataCollator()
# batch = collator(data)
# # input_ids 被 padding 对齐：
# # tensor([[101, 2003, 102,   0],
# #         [101, 3000, 4000, 102]])
# # labels 长度一致，直接堆叠：
# # tensor([0, 1])


'''
DefaultDataCollator

不会对任何的数据进行统一长度，需要输入的时候就统一长度，


'''


from transformers import DefaultDataCollator

# 正确的数据格式（labels 是单值）
data = [
    {"input_ids": [101, 2003, 102], "labels": 0},  # labels 是单值
    {"input_ids": [101, 3000, 4000], "labels": 1}  # labels 是单值
]

collator = DefaultDataCollator()
batch = collator(data)  # 正常运行，无报错
print(batch)


data = [
    {"input_ids": [101, 2003, 102], "labels": [0,0,0]},  # 长度 3
    {"input_ids": [101, 3000, 4000], "labels": [1,1,1]}  # 长度 4
]
collator = DefaultDataCollator()
batch = collator(data)  # 正常运行，无报错
print(batch)